{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generic Agent Architecture\n",
    "![Router](images/generic-agent-architecture.png)\n",
    "\n",
    "- pass tool message back to model and let the model to decide:\n",
    "    - make another call for a tool\n",
    "    - spond directly\n",
    "\n",
    "## ReAct (a general agent architecture)\n",
    "\n",
    "- **act** - take an action\n",
    "- **observe** - grab the response from the action\n",
    "- **reason** - analyze the response and determine what to do next"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## StockAnalysisAI\n",
    "\n",
    "Examples of user requests the agent should handle:\n",
    "\n",
    "- Should I invest in Tesla stocks?\n",
    "- How does Tesla compare to Apple?\n",
    "- What’s the stock symbol for Microsoft?\n",
    "- Fetch the latest Tesla stock data and analyze it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def lookup_stock_symbol(company_name: str) -> str:\n",
    "    \"\"\"\n",
    "    Converts a company name to its stock symbol using a financial API.\n",
    "\n",
    "    Parameters:\n",
    "        company_name (str): The full company name (e.g., 'Tesla').\n",
    "\n",
    "    Returns:\n",
    "        str: The stock symbol (e.g., 'TSLA') or an error message.\n",
    "    \"\"\"\n",
    "    api_url = \"https://www.alphavantage.co/query\"\n",
    "    params = {\n",
    "        \"function\": \"SYMBOL_SEARCH\",\n",
    "        \"keywords\": company_name,\n",
    "        \"apikey\": \"your_alphavantage_api_key\"\n",
    "    }\n",
    "    \n",
    "    response = requests.get(api_url, params=params)\n",
    "    data = response.json()\n",
    "    \n",
    "    if \"bestMatches\" in data and data[\"bestMatches\"]:\n",
    "        return data[\"bestMatches\"][0][\"1. symbol\"]\n",
    "    else:\n",
    "        return f\"Symbol not found for {company_name}.\"\n",
    "\n",
    "lookup_stock_symbol(\"Tesla\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yfinance as yf\n",
    "import json\n",
    "from pprint import pformat\n",
    "import ast\n",
    "\n",
    "def fetch_stock_data_raw(stock_symbol: str) -> dict:\n",
    "    \"\"\"\n",
    "    Fetches comprehensive stock data for a given symbol and returns it as a combined dictionary.\n",
    "\n",
    "    Parameters:\n",
    "        stock_symbol (str): The stock ticker symbol (e.g., 'TSLA').\n",
    "        period (str): The period to analyze (e.g., '1mo', '3mo', '1y').\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary combining general stock info and historical market data.\n",
    "    \"\"\"\n",
    "    period = \"1mo\"\n",
    "    try:\n",
    "        stock = yf.Ticker(stock_symbol)\n",
    "\n",
    "        # Retrieve general stock info and historical market data\n",
    "        stock_info = stock.info  # Basic company and stock data\n",
    "        stock_history = stock.history(period=period).to_dict()  # Historical OHLCV data\n",
    "\n",
    "        # Combine both into a single dictionary\n",
    "        combined_data = {\n",
    "            \"stock_symbol\": stock_symbol,\n",
    "            \"info\": stock_info,\n",
    "            \"history\": stock_history\n",
    "        }\n",
    "\n",
    "        return pformat(combined_data)\n",
    "\n",
    "    except Exception as e:\n",
    "        return {\"error\": f\"Error fetching stock data for {stock_symbol}: {str(e)}\"}\n",
    "\n",
    "data = fetch_stock_data_raw(\"TSLA\")\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Binding tools to the LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.tools import Tool\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Create tool bindings with additional attributes\n",
    "lookup_stock = Tool.from_function(\n",
    "    func=lookup_stock_symbol,\n",
    "    name=\"lookup_stock_symbol\",\n",
    "    description=\"Converts a company name to its stock symbol using a financial API.\",\n",
    "    return_direct=False  # Return result to be processed by LLM\n",
    ")\n",
    "\n",
    "fetch_stock = Tool.from_function(\n",
    "    func=fetch_stock_data_raw,\n",
    "    name=\"fetch_stock_data_raw\",\n",
    "    description=\"Fetches comprehensive stock data including general info and historical market data for a given stock symbol.\",\n",
    "    return_direct=False\n",
    ")\n",
    "\n",
    "toolbox = [lookup_stock, fetch_stock]\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "simple_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "llm_with_tools = simple_llm.bind_tools(toolbox)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Agent's node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import MessagesState\n",
    "\n",
    "# System message\n",
    "assistant_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a professional financial assistant specializing in stock market analysis and investment strategies. \n",
    "Your role is to analyze stock data and provide **clear, decisive recommendations** that users can act on, \n",
    "whether they already hold the stock or are considering investing.\n",
    "\n",
    "You have access to a set of tools that can provide the data you need to analyze stocks effectively. \n",
    "Use these tools to gather relevant information such as stock symbols, current prices, historical trends, \n",
    "and key financial indicators. Your goal is to leverage these resources efficiently to generate accurate, \n",
    "actionable insights for the user.\n",
    "\n",
    "Your responses should be:\n",
    "- **Concise and direct**, summarizing only the most critical insights.\n",
    "- **Actionable**, offering clear guidance on whether to buy, sell, hold, or wait for better opportunities.\n",
    "- **Context-aware**, considering both current holders and potential investors.\n",
    "- **Free of speculation**, relying solely on factual data and trends.\n",
    "\n",
    "### Response Format:\n",
    "1. **Recommendation:** Buy, Sell, Hold, or Wait.\n",
    "2. **Key Insights:** Highlight critical trends and market factors that influence the decision.\n",
    "3. **Suggested Next Steps:** What the user should do based on their current position.\n",
    "\n",
    "If the user does not specify whether they own the stock, provide recommendations for both potential buyers and current holders. Ensure your advice considers valuation, trends, and market sentiment.\n",
    "\n",
    "Your goal is to help users make informed financial decisions quickly and confidently.\n",
    "\"\"\"))\n",
    "\n",
    "# Node\n",
    "def assistant(state: MessagesState):\n",
    "   return {\"messages\": [llm_with_tools.invoke([assistant_system_message] + state[\"messages\"])]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import START, StateGraph\n",
    "from langgraph.prebuilt import tools_condition\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from IPython.display import Image, display\n",
    "\n",
    "# Graph\n",
    "builder = StateGraph(MessagesState)\n",
    "\n",
    "# Define nodes: these do the work\n",
    "builder.add_node(\"assistant\", assistant)\n",
    "builder.add_node(\"tools\", ToolNode(toolbox))\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"assistant\")\n",
    "builder.add_conditional_edges(\n",
    "    \"assistant\",\n",
    "    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
    "    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
    "    tools_condition,\n",
    ")\n",
    "builder.add_edge(\"tools\", \"assistant\")\n",
    "react_graph = builder.compile()\n",
    "\n",
    "# Show\n",
    "display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = react_graph.invoke({\"messages\": [HumanMessage(content=\"Should I invest in Tesla stocks?\")]})\n",
    "for message in messages['messages']:\n",
    "    message.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = react_graph.invoke({\"messages\": [HumanMessage(content=\"How does Tesla compare to Apple? What would be a better option to invest in?\")]})\n",
    "for message in messages['messages']:\n",
    "    message.pretty_print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
